# 0. 引入必要的包
from util import get, preprocess_image
from sklearn.model_selection import train_test_split
import numpy as np
import pickle
import os
from tqdm import tqdm
import glob
import time
import joblib

# 1. 读取配置文件中的信息
train_dir = get("train") # 获取 训练数据路径
char_styles = get("char_styles") # 获取 字符样式列表，注意: 必须是列标
new_size = get("new_size") # 获取 新图像大小元组, 注意: 必须包含h和w

# 2. 生成X,y 
print("# 读取训练数据并进行预处理，") 
# 获取所有的图像文件
image_files1 = glob.glob(r"E:\大三2\计算机视觉\shufa_app_CNN提取特征\shufa\train\train_篆*")
image_files2 = glob.glob(r"E:\大三2\计算机视觉\shufa_app_CNN提取特征\shufa\train\train_隶*")
image_files3 = glob.glob(r"E:\大三2\计算机视觉\shufa_app_CNN提取特征\shufa\train\train_草*")
image_files4 = glob.glob(r"E:\大三2\计算机视觉\shufa_app_CNN提取特征\shufa\train\train_行*")
image_files5 = glob.glob(r"E:\大三2\计算机视觉\shufa_app_CNN提取特征\shufa\train\train_楷*")

# 将所有的图片文件放入一个列表中
image_files_list = [image_files1, image_files2, image_files3, image_files4, image_files5]

# 对每个文件调用 preprocess_image 函数
X = []
y = []
for i,image_files in enumerate(image_files_list):
    for file_name in tqdm(image_files,desc=f"正在处理 {char_styles[i]} 字体的图片"):
        # 检查 file_name 是否是一个图像文件
        if file_name.endswith('.jpg') or file_name.endswith('.png'):
            # 如果是，处理图像并将结果添加到 X
            X.append(preprocess_image(file_name, new_size))
            # 提取标签并添加到 y
            label = file_name.split('_')[1]
            # 将样式的下标添加到 y
            y.append(i)

# 将列表转换为 NumPy 数组
X = np.array(X)
y = np.array(y)

# 3. 分割测试集和训练集
print("# 将数据按 80% 和 20% 的比例分割")
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 4. 打印样本维度和类型信息
print("X_train: ", X_train.shape, X_train.dtype)  # 训练集特征的维度和类型
print("X_test: ", X_test.shape, X_test.dtype)  # 测试集特征的维度和类型
print("y_train: ", y_train.shape, y_train.dtype)  # 训练集标签的维度和类型
print("y_test: ", y_test.shape, y_test.dtype)  # 测试集标签的维度和类型

# 5. 序列化分割后的训练和测试样本
# 获取开始时间
start_time = time.time()

# 序列化分割后的训练和测试样本
file_path = './Xys/Xy'
joblib.dump((X_train, X_test, y_train, y_test), file_path)

# 获取结束时间
end_time = time.time()

# 计算运行时间
run_time = end_time - start_time

# 获取文件大小并转换为兆字节（MB）
file_size_MB = os.path.getsize(file_path) / (1024 * 1024)
print(" 保存 X_train, X_test, y_train, y_test 到./Xys/Xy.pkl 文件中")
print(f"文件已保存到 {file_path}")
print(f"文件大小为 {file_size_MB:.2f} MB")
print(f"保存文件运行时间为 {run_time} 秒")


